import subprocess
import csv
import os
from datetime import datetime
from openpyxl import Workbook
import random
import re


def extract_five_tuple_with_tshark(pcap_file_path):
    """
    使用tshark从pcap文件中提取五元组信息和HTTP方法
    支持上行和下行流量过滤
    """
    # tshark路径
    TSHARK_PATH = r"D:\Program Files\Wireshark\tshark.exe"

    # 检查tshark是否存在
    if not os.path.exists(TSHARK_PATH):
        print(f"tshark未找到: {TSHARK_PATH}")
        return []

    # 保留基本的五元组和HTTP方法字段，添加HTTP URI路径
    cmd = [
        TSHARK_PATH,
        '-r', pcap_file_path,
        '-Y', 'http && ip.addr == 192.168.5.241',
        '-T', 'fields',
        '-e', 'frame.time',
        '-e', 'ip.src',
        '-e', 'tcp.srcport',
        '-e', 'ip.dst',
        '-e', 'tcp.dstport',
        '-e', 'ip.proto',
        '-e', 'http.host',
        '-e', 'http.request.method',
        '-e', 'http.request.uri',
        '-e', 'http.request.version',  # 添加HTTP版本信息
        '-e', 'http.response.code',
        '-e', 'http.response.phrase',
        '-e', 'tcp.stream'
    ]

    try:
        # 执行tshark命令
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            encoding='utf-8',
            shell=True
        )

        if result.returncode != 0:
            print(f"tshark执行出错: {result.stderr}")
            return []

        # 解析输出结果
        lines = result.stdout.strip().split('\n')
        results = []
        cyber_stream_ids = set()

        # 第一遍扫描：找出包含"cyber"的请求及其TCP流ID
        for line in lines:
            if line.strip():
                fields = line.split('\t')
                if len(fields) >= 12:  # 现在需要至少12个字段
                    http_host = fields[6] if len(fields) > 6 else ''
                    tcp_stream_id = fields[12] if len(fields) > 12 else ''
                    
                    # 如果host包含cyber，则记录其TCP流ID
                    if 'cyber' in http_host.lower():
                        cyber_stream_ids.add(tcp_stream_id)

        # 第二遍扫描：收集所有与cyber请求相关的数据包
        for line in lines:
            if line.strip():
                fields = line.split('\t')
                if len(fields) >= 12:  # 现在需要至少12个字段
                    # 构造完整的HTTP请求信息（类似Wireshark的info列）
                    http_method = fields[7] if len(fields) > 7 else ''
                    http_uri = fields[8] if len(fields) > 8 else ''
                    http_version = fields[9] if len(fields) > 9 else 'HTTP/1.1'
                    http_response_code = fields[10] if len(fields) > 10 else ''
                    http_response_phrase = fields[11] if len(fields) > 11 else ''
                    tcp_stream_id = fields[12] if len(fields) > 12 else ''
                    http_host = fields[6] if len(fields) > 6 else ''
                    
                    # 只处理与cyber请求相关的流
                    if tcp_stream_id in cyber_stream_ids:
                        # 构造完整的HTTP请求/响应行
                        full_http_info = ''
                        if http_method and http_uri:
                            full_http_info = f"{http_method} {http_uri} {http_version}"
                        elif http_response_code and http_response_phrase:
                            # 如果是响应，则构造响应行
                            full_http_info = f"{http_version} {http_response_code} {http_response_phrase}"
                        
                        result_item = {
                            'timestamp': fields[0] if len(fields) > 0 else '',
                            'sip': fields[1] if len(fields) > 1 else '',
                            'spt': fields[2] if len(fields) > 2 else '',
                            'dip': fields[3] if len(fields) > 3 else '',
                            'dpt': fields[4] if len(fields) > 4 else '',
                            'proto': fields[5] if len(fields) > 5 else '',
                            'http_host': http_host,
                            'http_method': http_method,
                            'full_http_info': full_http_info,  # 完整的HTTP请求/响应信息
                            'tcp_stream_id': tcp_stream_id  # TCP流ID，用于关联请求和响应
                        }
                        results.append(result_item)

        return results

    except Exception as e:
        print(f"执行tshark时出错: {e}")
        import traceback
        traceback.print_exc()
        return []


def save_to_xlsx(data, output_file):
    """
    将提取的数据保存到Excel文件
    """
    if not data:
        print("没有数据需要保存")
        return

    # 创建工作簿
    wb = Workbook()
    ws = wb.active
    ws.title = "Five Tuple Data"

    # 保留基本字段的表头，只保留完整的HTTP信息字段，不单独显示uri和version
    fieldnames = ['timestamp', 'sip', 'spt', 'dip', 'dpt', 'proto', 'http_host', 'http_method', 'full_http_info', 'tcp_stream_id']
    ws.append(fieldnames)

    # 写入数据
    for item in data:
        row = [item[field] for field in fieldnames]
        ws.append(row)

    # 保存文件
    wb.save(output_file)
    print(f"Excel文件已保存到: {output_file}")


def main():
    # 使用正则表达式查找并处理目录中的所有数字命名的pcap文件
    data_dir = r"D:\DPI测试\转储测试\Data"
    
    # 查找目录中所有符合数字命名规则的pcap文件
    pattern = re.compile(r'^(\d+)\.pcap$')
    pcap_files = [f for f in os.listdir(data_dir) if pattern.match(f)]
    
    if not pcap_files:
        print(f"在目录 {data_dir} 中未找到数字命名的pcap文件")
        return
    
    # 按数字顺序排序文件
    pcap_files.sort(key=lambda x: int(pattern.match(x).group(1)))
    
    print(f"找到以下数字命名的pcap文件: {pcap_files}")
    
    # 指定输出目录
    output_dir = r"D:\DPI测试\转储测试\Data"
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 处理每个找到的文件
    for filename in pcap_files:
        pcap_file_path = os.path.join(data_dir, filename)
        print(f"\n正在处理文件: {pcap_file_path}")
        
        # 检查输入文件是否存在
        if os.path.exists(pcap_file_path):
            # 提取数据
            results = extract_five_tuple_with_tshark(pcap_file_path)
            
            if results:
                # 生成输出文件名，使用Wire前缀，后跟日期和随机数字（根据规范不使用下划线）
                current_date = datetime.now().strftime('%m%d')
                random_number = random.randint(1000, 9999)
                output_filename = f"Wire{current_date}{random_number}.xlsx"
                output_file_path = os.path.join(output_dir, output_filename)
                
                # 保存到Excel文件
                save_to_xlsx(results, output_file_path)
                print(f"提取完成，共找到 {len(results)} 条记录")
                
                # 显示前几条记录
                print("\n前5条记录:")
                for i, record in enumerate(results[:5]):
                    print(f"{i + 1}. {record['sip']}:{record['spt']} -> {record['dip']}:{record['dpt']} ")
                    print(f"   METHOD:{record['http_method']} HOST:{record['http_host']}")
                    print(f"   FULL INFO:{record['full_http_info']} TCP STREAM:{record['tcp_stream_id']}")
            else:
                print("未找到符合条件的记录")
        else:
            print(f"输入文件不存在: {pcap_file_path}")


if __name__ == "__main__":
    main()